package com.tezign.intelligence.mini.dam.aifamilyconsultant.util;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import lombok.AllArgsConstructor;
import lombok.Data;
import org.apache.http.HttpEntity;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.util.EntityUtils;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.io.IOException;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;

public class ChatGPTUtil {

    private static final String AI_API_URL = "https://api.302.ai/v1/chat/completions"; // 替换为实际302.ai API地址
    private static final String API_KEY = "sk-DxuqVi2JpSQ7brhuDL4bD7KSvoYT89D9FJkHdJSl3a3w9FpL"; // 替换为你的API密钥

    public static Flux<String> streamChat(List<ChatMessage> messages, String model) {
        WebClient client = WebClient.create(AI_API_URL);
        // 构建302.ai API请求chatgpt-4o-latest
        AiChatRequest request = new AiChatRequest(
                model, // 根据302.ai文档使用正确的模型名称
                messages,
                true, // 启用流式响应
                0.7 // 温度参数
        );

        // 调用302.ai API并处理流式响应
        return client.post()
                .header("Authorization", "Bearer " + API_KEY)
                .header("Content-Type", "application/json")
                .header("Accept", "text/event-stream")
                .bodyValue(request)
                .retrieve()
                .bodyToFlux(String.class);
    }

    // 解析302.ai API返回的chunk数据
    public static String parseChunk(String chunk) {
        try {
            // 这里需要根据302.ai API实际的流式响应格式进行解析
            // 示例实现，实际需要调整
            JsonNode node = new ObjectMapper().readTree(chunk);
            if (node.has("choices") && node.get("choices").size() > 0) {
                JsonNode choice = node.get("choices").get(0);
                if (choice.has("delta") && choice.get("delta").has("content")) {
                    return choice.get("delta").get("content").asText();
                }
            }
        } catch (Exception e) {
            return "";
        }
        return "";
    }

    public static String chatCompletion(String role, String content, String model) {
        WebClient client = WebClient.create(AI_API_URL);
        // 构建302.ai API请求

        final List<ChatMessage> messages = new ArrayList<>();
        final ChatMessage systemMessage = ChatMessage.builder().role(ChatMessageRole.SYSTEM).content(role).build();
        final ChatMessage userMessage = ChatMessage.builder().role(ChatMessageRole.USER).content(content).build();
        messages.add(systemMessage);
        messages.add(userMessage);
        AiChatRequest request = new AiChatRequest(
                model, // 根据302.ai文档使用正确的模型名称
                messages,
                false, // 启用流式响应
                0.7 // 温度参数
        );
        Mono<String> stringMono = client.post()
                .header("Authorization", "Bearer " + API_KEY)
                .header("Content-Type", "application/json")
                .header("Accept", "text/event-stream")
                .bodyValue(request)
                .retrieve()
                .bodyToMono(String.class)
                .onErrorResume(e -> {
                    // 错误处理
                    return Mono.just("Error: " + e.getMessage());
                });
        String block = stringMono.block(Duration.ofSeconds(30));
        return parseReportChunk(block);
    }

    // 解析302.ai API返回的chunk数据
    public static String parseReportChunk(String chunk) {
        try {
            // 这里需要根据302.ai API实际的流式响应格式进行解析
            // 示例实现，实际需要调整
            JsonNode node = new ObjectMapper().readTree(chunk);
            if (node.has("choices") && node.get("choices").size() > 0) {
                JsonNode choice = node.get("choices").get(0);
                if (choice.has("message") && choice.get("message").has("content")) {
                    return choice.get("message").get("content").asText();
                }
            }
        } catch (Exception e) {
            return "";
        }
        return "";
    }

    // 内部类定义请求和消息结构
    @Data
    @AllArgsConstructor
    private static class AiChatRequest {
        private String model;
        private List<ChatMessage> messages;
        private boolean stream;
        private double temperature;
    }
}